import torch
import networkx as nx
import matplotlib.pyplot as plt
# 数据集 Zachary's karate club network
# 一个俱乐部中有 34个会员，之间的互相认识的关系有 78 条 ，管理人员和教练之间出现冲突，34个人站队，4中最终选择。
# https://pytorch-geometric.readthedocs.io/en/latest/modules/datasets.html
from torch_geometric.datasets import KarateClub

dataset = KarateClub()
print("Dataset:", dataset)
print("图长度:", len(dataset))
print("图节点数:", dataset.num_features)
print("图分类数:", dataset.num_classes)

data = dataset[0]
# Data(x=[34, 34], edge_index=[2, 156], y=[34], train_mask=[34])
print("train_mask",data.train_mask)
"""
1, 1, 1, 1, 3, 3, 3, 1, 0, 1, 3, 1, 1, 1, 0, 0, 3, 1, 0, 1, 0, 1, 0, 0, 2, 2, 0, 0, 2, 0, 0, 2, 0, 0
T, F, F, F, T, F, F, F, T, F, F, F, F, F, F, F, F, F, F, F, F, F, F, F, T, F, F, F, F, F, F, F, F, F

x 有34个点，每个的特征是一个34维的向量。
edge_index 邻接矩阵关系的映射向量，由于是关系表示因此每个点和另外一个点有关系就会建立一条关系，

假设图中有2条边，每条边都建立一个 [2,2] 的关系 
[[A,B],
[B,A],
[C,D],
[D,C]] (4 行 2 列) 由于有78条边 ，所以应该建立一个 （156行2列的关系，这里是index的转置）
y 是标签
train_mask 屏蔽了没有标签的数据，不参与loss计算
"""
# print(data)
# 稀疏矩阵来表示邻接矩阵
edge_index = data.edge_index
print(edge_index.t())

# 使用networksx进行可视化展示
from torch_geometric.utils import to_networkx
G = to_networkx(data, to_undirected=True)
from Visualize import visualize_graph
visualize_graph(G, color=data.y)

